package dongshi.daddy.trieTree;

import java.util.HashSet;
import java.util.Set;

public class TrieTree {

    public TreeNode root;

    public TrieTree() {
        root = new TreeNode();
    }


    /**
     * 在前缀树中插入字符串
     * 这种++的方法，导致，一个node，有多少个end，就有多少个相同的字符串
     * 一个node，有多少个path，就有多少个字符串经过（root的path代表有多少个字符串）（字符串末尾的node的path也会++）
     *
     * @param string 被插入的字符串（以前插入过的也可以插入）
     */
    public void insertString(String string) {
        if (string == null || string.length() == 0) {
            return;
        }
        int length = string.length();
        TreeNode nowNode = root;
        for (int i = 0; i < length; i++) {
            char now = string.charAt(i);
            int index = now - 'a';
            //index为字符now所处的位置
            if (nowNode.next[index] == null) {
                nowNode.next[index] = new TreeNode();
            }
            nowNode.isLeaf = false;
            // 先对当前node的path++，再转移到下一个node
            nowNode.path++;
            nowNode = nowNode.next[index];
        }
        // 处理 ab abc ，通过前缀a查询，也需要查询出ab的情况
        nowNode.isWordEnd = true;
        //在最后的node，path和end++
        nowNode.path++;
        nowNode.end++;
    }

    /**
     * 返回这个前缀树总共插入了多少个字符串
     *
     * @return
     */
    public int size() {
        return root.path;
    }

    /**
     * 前缀树查询总共插入这个字符串多少次，如果没插入过，则返回0
     *
     * @param string
     * @return
     */
    public int getStringNum(String string) {
        if (string == null || string.length() == 0) {
            return 0;
        }
        int length = string.length();
        TreeNode nowNode = root;
        for (int i = 0; i < length; i++) {
            char now = string.charAt(i);
            int index = now - 'a';
            //如果没有这个节点，说明不存在，直接返回0
            if (nowNode.next[index] == null) {
                return 0;
            }
            nowNode = nowNode.next[index];
        }
        //此时nowNode已经处于最后一个节点
        return nowNode.end;
    }

    /**
     * 前缀树查询以这个字符串为前缀的字符串总共多少个（包括以他为结尾的）
     *
     * @param string 前缀
     * @return
     */
    public int getPrefixNum(String string) {
        if (string == null || string.length() == 0) {
            return 0;
        }
        int length = string.length();
        TreeNode nowNode = root;
        for (int i = 0; i < length; i++) {
            char now = string.charAt(i);
            int index = now - 'a';
            //如果没有这个节点，说明前缀不存在，直接返回0
            if (nowNode.next[index] == null) {
                return 0;
            }
            nowNode = nowNode.next[index];
        }
        //此时nowNode已经处于前缀的最后一个节点
        return nowNode.path;
    }

//    public List<String> findByPrefix(String prefix) {
    public Set<String> findByPrefix(String prefix) {
        // 注意：根节点不存储任何元素
        TreeNode curNode = root;

        int prefixLen = prefix.length();
        // 1：找到prefix对应的TreeNode对象
        for (int i = 0; i < prefixLen; i++) {
            int idx = prefix.charAt(i) - 'a';
            TreeNode[] dataArr = curNode.next;
            if (dataArr[idx] == null) {
                System.out.println("not find!");
                return null;
            }
            // 非前缀的最后一个元素，遇到空，则说明要匹配的前缀不存在
            /*if (dataArr[idx] != null) {
                if (i > prefixLen - 1) {
                    return null;
                } else {
                    curNode = dataArr[idx];
                }
            }*/
            // 继续向下
            curNode = dataArr[idx];
        }
        // 2：根据prefix对应的TreeNode对象，递归找到所有的可能字符串
        TreeNode[] possibleTreeNodeArr = curNode.next;
        // 3：递归找到所有的可能字符串
//        List<String> possibleStrList = new ArrayList<>();
        Set<String> possibleStrList = new HashSet<>();
        /*for (int i = 0; i < possibleTreeNodeArr.length; i++) {
            if (possibleTreeNodeArr[i] != null) possibleStrList.add(prefix + (char) (i + 'a'));
        }
        for (int i = 0; i < possibleTreeNodeArr.length; i++) {
            queryAllPossibleStr(i, possibleTreeNodeArr, possibleStrList, prefix);
        }*/
        queryAllPossibleStr(0, curNode, possibleTreeNodeArr, possibleStrList, prefix);
        return possibleStrList;
    }

    private void queryAllPossibleStr(int i, TreeNode curNode, TreeNode[] possibleTreeNodeArr, Set<String> possibleStrList, String prefix) {
        if (i >= possibleTreeNodeArr.length || possibleTreeNodeArr == null) return;
        String newPrefix = prefix + (char) (i + 'a');
        // 元素为null，说明到达叶子节点
        if ((possibleTreeNodeArr[i] == null && curNode.isLeaf) || curNode.isWordEnd) {
//        if (possibleTreeNodeArr[i] == null && i == possibleTreeNodeArr.length - 1) {
//        if (possibleTreeNodeArr[i] != null && possibleTreeNodeArr[i].isLeaf) {
            possibleStrList.add(prefix);
            // 下层
//            queryAllPossibleStr(0, possibleTreeNodeArr[i], possibleTreeNodeArr[i].next, possibleStrList, newPrefix);
        } /*else {
            // 当前无元素，则向右继续找，有则向下和向右找
            if (possibleTreeNodeArr[i] != null) {
                // 下层
                queryAllPossibleStr(0, possibleTreeNodeArr[i], possibleTreeNodeArr[i].next, possibleStrList, newPrefix);
            }
        }*/
        // 当前无元素，则向右继续找，有则向下和向右找
        if (possibleTreeNodeArr[i] != null) {
            // 下层
            queryAllPossibleStr(0, possibleTreeNodeArr[i], possibleTreeNodeArr[i].next, possibleStrList, newPrefix);
        }
        // 不管咋的，都得向右→
        queryAllPossibleStr(i + 1, curNode, possibleTreeNodeArr, possibleStrList, prefix);
    }
}
